package com.smh.aiagent.agent;

import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.smh.aiagent.agent.model.AgentState;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.DefaultToolCallingManager;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.stream.Collectors;

/**
 * @author shiminghui
 * @date 2025/6/8 15:37
 * @description: TODO
 */
@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public abstract class ToolCallingAgent extends ReactAgent {

    private final ToolCallback[] availableTools;

    private final ToolCallingManager toolCallingManager;
    private final ChatOptions toolCallingOptions;
    private ChatResponse toolCallingResponse;


    public ToolCallingAgent(ToolCallback[] allTools) {
        super();
        this.availableTools = allTools;
        this.toolCallingManager = DefaultToolCallingManager.builder().build();
        this.toolCallingOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true) // 将工具调用代理给工具执行器,SpringAi不再进行处理
                .build();
    }

    @Override
    protected boolean think() {
        if (getNextStepPrompt() != null && !getNextStepPrompt().isEmpty()) {
            getMessageList().add(new UserMessage(getNextStepPrompt()));
        }
        // 调用AI模型
        Prompt prompt = new Prompt(getMessageList(), this.toolCallingOptions);

        try {
            ChatResponse chatResponse = getChatClient().prompt(prompt)
                    .system(getSystemPrompt())
                    .tools(availableTools)
                    .call()
                    .chatResponse();
            // 记录响应结果,用于后续action使用
            this.toolCallingResponse = chatResponse;
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();

            String result = assistantMessage.getText();
            log.info("{}的思考结果为:{}", getName(), result);

            // 解析ai的结果,判断是否要进行工具调用
            if (!chatResponse.hasToolCalls()) {
                // 仅在没有工具调用时才添加结果,因为存在工具调用时,工具调用会自动将AI结果添加到messageList中,这样会造成重复
                getMessageList().add(assistantMessage);
                return false;
            }

            String toolCallingContent = assistantMessage.getToolCalls().stream()
                    .map(toolCall -> String.format("工具名称:[%s],参数:[%s]", toolCall.name(), toolCall.arguments()))
                    .collect(Collectors.joining("\n"));
            log.info("{}使用了以下工具信息:\n{}", getName(), toolCallingContent);

            return true;
        } catch (Exception e) {
            log.error("{}思考时遇到问题:{}", getName(), e.getMessage());
            getMessageList().add(new AssistantMessage("处理时遇到问题:" + e.getMessage()));
            return false;
        }
    }

    @Override
    protected String action() {
        // 工具调用
        if (!toolCallingResponse.hasToolCalls()) {
            return "没有工具需要调用";
        }
        Prompt prompt = new Prompt(getMessageList(), toolCallingOptions);
        try {
            ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallingResponse);

            // 记录消息上下文
            setMessageList(toolExecutionResult.conversationHistory());

            // 返回结果
            Message lastMessage = toolExecutionResult.conversationHistory().getLast();
            if (lastMessage instanceof ToolResponseMessage toolResponseMessage) {
                String toolCallingResult = toolResponseMessage.getResponses().stream().map(response -> {
                            if ("doTerminate".equals(response.name())) {
                                setAgentState(AgentState.FINISHED);
                            }
                            return String.format("%s工具执行完成,执行结果为:%s", response.name(), response.responseData());
                        })
                        .collect(Collectors.joining("\n"));
                log.info("{}的工具调用结果为:{}", getName(), toolCallingResult);
                return toolCallingResult;
            } else {
                return "没有工具调用结果";
            }
        } catch (Exception e) {
            log.error("{}工具调用时遇到问题:{}", getName(), e.getMessage());
            return "工具调用时遇到问题:" + e.getMessage();
        }
    }
}
